//
// Created by Panfeng Guo on 2022/11/12.
//
#include <iostream>
#include <string>

#include "common/exception.h"
#include "common/logger.h"
#include "common/rid.h"
#include "storage/index/b_plus_tree.h"
#include "storage/page/header_page.h"

template <typename KeyType, typename ValueType, typename KeyComparator>
BPlusTree<KeyType, ValueType, KeyComparator>::
BPlusTree(const std::string &name, BufferPoolManager *buffer_pool_manager, const KeyComparator &comparator,
          page_id_t root_page_id)
        :index_name_(name), root_page_id_(root_page_id),
         buffer_pool_manager_(buffer_pool_manager), comparator_(comparator){}

template <typename KeyType, typename ValueType, typename KeyComparator>
bool BPlusTree<KeyType, ValueType, KeyComparator>::
IsEmpty() const {
    return root_page_id_ == INVALID_PAGE_ID;
}

/*!
 * Insert: insert a kv into b+tree
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @param key
 * @param value
 * @param transaction
 * @return: if insertion is successful, return true; otherwise, return false
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
bool BPlusTree<KeyType, ValueType, KeyComparator>::
Insert(const KeyType &key, const ValueType &value, Transaction *transaction) {
    std::lock_guard<std::mutex> lock(mutex_);
    if(IsEmpty()){ // case 1: there is no tree, create one
        // step 1: create a page
        auto page = buffer_pool_manager_->NewPage(&root_page_id_);
        if(page != nullptr){
            throw Exception(ExceptionType::INDEX_TYPE,
                            "all pages are pinned, can not create a new tree");
        }
        // step 2: update root_page_id_ in header page
        UpdateRootPageId(true);
        // step 3: insert kv into root page
        auto root = reinterpret_cast<BPlusTreeLeafPage<KeyType, ValueType, KeyComparator> *> (page->GetData());
        root->Init(root_page_id_);
        root->Insert(key, value, comparator_);
        buffer_pool_manager_->UnpinPage(root_page_id_, true);
        return true;
    }else{ // case 2: insert into leaf
        return InsertIntoLeaf(key, value, transaction);
    }
}


/*!
 * UpdateRootPageId: update root_page_id_ in header page
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @param insert_record : whether it it a insert operation
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
void BPlusTree<KeyType, ValueType, KeyComparator>::
UpdateRootPageId(bool insert_record){
    auto *page = buffer_pool_manager_->FetchPage(HEADER_PAGE_ID);
    if (page == nullptr){
        throw Exception(ExceptionType::INDEX_TYPE, "all page are pinned, can not update page");
    }
    auto *header_page = reinterpret_cast<HeaderPage *>(page->GetData());

    if(insert_record){
        // insert a new record in header_page
        header_page->InsertRecord(index_name_, root_page_id_);
    }else{
        // update root_page_id_ in header_page
        header_page->UpdateRecord(index_name_, root_page_id_);
    }
    buffer_pool_manager_->UnpinPage(HEADER_PAGE_ID, true);
}

template <typename KeyType, typename ValueType, typename KeyComparator>
bool BPlusTree<KeyType, ValueType, KeyComparator>::
InsertIntoLeaf(const KeyType &key, const ValueType &value, Transaction *transaction){
    // step 1: find the leaf to insert
    auto leaf = FindLeafPage(key, false, Operation::INSERT, transaction);
    if(leaf == nullptr){
        return false;
    }

    // step 2: if key already exist, return false
    ValueType v;
    if(leaf->Lookup(key, v, comparator_)){
        UnlockUnpinPages(Operation::INSERT, transaction);
        return false;
    }

    // step 3: insert into leaf
    if(leaf->GetSize() < leaf->GetMaxSize()){ // case 1: insert without splitting
        leaf->Insert(key, value, comparator_);
    }else{ // case 2: leaf page is full, need to split
        // create a new leaf page
        auto new_leaf = Split<BPlusTreeLeafPage<KeyType, ValueType, KeyComparator>>(leaf);
        new_leaf->SetNextPageId(leaf->GetNextPageId());
        leaf->SetNextPageId(new_leaf->GetPageId());
        // insert into page
        if(comparator_(key, new_leaf->KeyAt(0)) < 0){
            leaf->Insert(key, value, comparator_);
        } else{
            new_leaf->Insert(key, value, comparator_);
        }
        // insert into parent
        InsertIntoParent(leaf, new_leaf.KeyAt(0), new_leaf, transaction);
    }
    UnlockUnpinPages(Operation::INSERT, transaction);
    return true;
}

/*!
 * FindLeafPage: find leaf page of the given key
 * all node passing by must be locked and pinned, unless deemed safe
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @param key
 * @param leftMost: TODO: fill this
 * @param op
 * @param transaction: can be nullptr when readonly
 * @return: the leaf page which contains the given key
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
BPlusTreeLeafPage<KeyType, ValueType, KeyComparator> *
BPlusTree<KeyType, ValueType, KeyComparator>::
FindLeafPage(const KeyType &key, bool leftMost, Operation op, Transaction *transaction){
    // step 1: is operation is not read, lock root
    if(op != Operation::READONLY){
        lockRoot();
        root_is_locked = true;
    }

    // step 2: get root page
    if(IsEmpty()){
        return false;
    }
    auto root_page = buffer_pool_manager_->FetchPage(root_page_id_);
    if(root_page == nullptr){
        throw Exception(ExceptionType::INDEX_TYPE, "all page are pinned, can not update page");
    }

    // step 3: lock root page and add into transaction page set
    if(op == Operation::READONLY){
        root_page->RLatch();
    }else{
        root_page->WLatch();
    }
    if(transaction != nullptr){
        transaction->AddIntoPageSet(root_page);
    }

    // step 4: traverse b plus tree
    auto node = reinterpret_cast<BPlusTreePage *>(root_page->GetData());
    auto parent_page = root_page;
    while(!node->IsLeafPage()){
        page_id_t parent_page_id, child_page_id;
        // convert to internal page
        auto internal_page = reinterpret_cast<BPlusTreeInternalPage<KeyType, ValueType, KeyComparator> *>(node);
        parent_page_id = internal_page->GetParentPageId();
        if(leftMost){
            child_page_id = internal_page->ValueAt(0);
        }else{
            child_page_id = internal_page->Lookup(key, comparator_);
        }
        // get child page
        auto child_page = buffer_pool_manager_->FetchPage(child_page_id);
        if(child_page == nullptr){
            throw Exception(ExceptionType::INDEX_TYPE, "all page are pinned, can not update page");
        }

        // lock child page and unlock parent page
        if(op == Operation::READONLY){
            child_page->RLatch();
            UnlockUnpinPages(op, transaction);
        }else{
            child_page->WLatch();
            if(isSafe(node, op)){
                // for insertion and deletion, if node won't aggregate or split, unlock all parent
                UnlockUnpinPages(op, transaction);
            }
        }

        // add child page into transaction page set
        if(transaction != nullptr){
            transaction->AddIntoPageSet(child_page);
        }else{ // is transaction is nullptr, unlock and unpin parent
            parent_page->RUnlatch();
            buffer_pool_manager_->UnpinPage(parent_page->GetPageId(), false);
            parent_page = child_page;
        }
    }
    auto leaf_page = reinterpret_cast<BPlusTreeLeafPage<KeyType, ValueType, KeyComparator>*>(node);
    return leaf_page;
}

/*!
 * isSafe: judge whether a b plus tree node is safe
 * for insertion: if node.size() < max_size, so it must be safe
 * for deletion: if node.size() > min_size, so it must be safe
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @tparam N
 * @param node
 * @param op
 * @return: safe or not
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
template <typename N>
bool BPlusTree<KeyType, ValueType, KeyComparator>::
isSafe(N *node, Operation op)
{
    if (op == Operation::INSERT)
    {
        return node->GetSize() < node->GetMaxSize();
    }else if (op == Operation::DELETE)
    {
        // >=: keep same with `coalesce logic`
        return node->GetSize() > node->GetMinSize() + 1;
    }
    return true;
}

/*!
 * UnlockUnpinPages: unlock and unpin all transaction visited page
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @param op
 * @param transaction
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
void BPlusTree<KeyType, ValueType, KeyComparator>::
UnlockUnpinPages(Operation op, Transaction *transaction){
    if(transaction == nullptr){
        return;
    }
    // step 1: unlock and unpin all transaction page
    for(auto page : *(transaction->GetPageSet())){
        if(op == Operation::READONLY){
            page->RUnlatch();
            buffer_pool_manager_->UnpinPage(page->GetPageId(), false);
        }else{
            page->WUnlatch();
            buffer_pool_manager_->UnpinPage(page->GetPageId(), true);
        }
    }
    transaction->GetPageSet()->clear();

    // step 2: delete page
    for(auto page_id : *(transaction->GetDeletedPageSet())){
        buffer_pool_manager_->DeletePage(page_id);
    }

    // step 3: unlock root
    if(root_is_locked){
        root_is_locked = false;
        unlockRoot();
    }
}

/*!
 * Split: split a node into two
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @tparam N
 * @param node: the node to split
 * @return: the new node generated after splitting
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
template <typename N>
N *BPlusTree<KeyType, ValueType, KeyComparator>::
Split(N *node){
    page_id_t new_page_id;
    auto new_page = buffer_pool_manager_->NewPage(&new_page_id);
    if(new_page == nullptr){
        throw Exception(ExceptionType::INDEX_TYPE, "all page are pinned, can not split page");
    }
    auto new_node = reinterpret_cast<N *>(new_page->GetData());
    new_node->Init(new_page_id);
    node->MoveHalfTo(new_node, buffer_pool_manager_);
    return new_node;
}

/*!
 * InsertIntoParent:
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @param old_node
 * @param key
 * @param new_node
 * @param transaction
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
void BPlusTree<KeyType, ValueType, KeyComparator>::
InsertIntoParent(BPlusTreePage *old_node, const KeyType &key,
                 BPlusTreePage *new_node, Transaction *transaction){
    if(old_node->IsRootPage()){ // case 1: old node is root, generate a new root
        // generate new root page
        auto page = buffer_pool_manager_->NewPage(&root_page_id_);
        if(page == nullptr){
            throw Exception(ExceptionType::INDEX_TYPE,
                            "all page are pinned, can not insert into parent");
        }
        assert(page->GetPinCount() == 1);
        auto root = reinterpret_cast<BPlusTreeInternalPage<KeyType, ValueType, KeyComparator> *>(page->GetData());
        root->Init(root_page_id_);
        // link new root page and two child
        root->PopulateNewRoot(old_node->GetPageId(), key, new_node->GetPageId());
        old_node->SetParentPageId(root_page_id_);
        new_node->SetParentPageId(root_page_id_);
        // update root page in header page
        UpdateRootPageId(false);

        buffer_pool_manager_->UnpinPage(new_node->GetPageId(), true);
        buffer_pool_manager_->UnpinPage(old_node->GetPageId(), true);
    }else{ // case 2: old node is not root
        auto page = buffer_pool_manager_->FetchPage(old_node->GetParentPageId());
        if(page == nullptr){
            throw Exception(ExceptionType::INDEX_TYPE,
                            "all page are pinned, can not insert into parent");
        }
        auto parent = reinterpret_cast<BPlusTreeInternalPage<KeyType, ValueType, KeyComparator> *>(page->GetData());
        if (parent->GetSize() < parent->GetMaxSize()){ // case 1: parent have free space
            parent->InsertNodeAfter(old_node->GetPageId(), key, new_node->GetPageId());
            new_node->SetParentPageId(parent->GetPageId());
            buffer_pool_manager_->UnpinPage(new_node->GetPageId(), true);
        }else{ // case 2: parent have no free space
            page_id_t  tmp_page_id;
            auto tmp_page = buffer_pool_manager_->NewPage(&tmp_page_id);
            if(tmp_page == nullptr){
                throw Exception(ExceptionType::INDEX_TYPE,
                                "all page are pinned, can not insert into parent");
            }
            assert(tmp_page->GetPinCount() == 1);
            auto tmp_node = reinterpret_cast<BPlusTreeInternalPage<KeyType, ValueType, KeyComparator> *>(tmp_page->GetData());
            tmp_node->Init(tmp_page_id);
            tmp_node->SetSize(parent->GetSize());
            // move 1~n and new node to tmp_node
            for(int i = 1, tmp_idx; i < parent->GetSize() && tmp_idx < tmp_node->GetSize(); i++, tmp_idx++){
                tmp_node->SetKeyAt(tmp_idx, parent->KeyAt(i));
                tmp_node->SetValueAt(tmp_idx, parent->ValueAt(i));
                if(parent->ValueAt(i) == old_node->GetPageId()){ // insert new node
                    tmp_idx++;
                    tmp_node->SetKeyAt(tmp_idx, key);
                    tmp_node->SetValueAt(tmp_idx, new_node->GetPageId());
                }
            }
            assert(tmp_node->GetSize() == tmp_node->GetMaxSize());
            // split half of tmp_node to new_parent
            auto new_parent = Split<BPlusTreeInternalPage<KeyType, ValueType, KeyComparator>>(tmp_node);
            // move tmp_node to parent
            parent->SetSize(tmp_node->GetSize()+1);
            for(int i = 0; i < parent->GetSize(); i++){
                parent->SetKeyAt(i+1, tmp_node->KeyAt(i));
                parent->SetValueAt(i+1, tmp_node->ValueAt(i));
            }
            // set child node's parent
            if(comparator_(key, new_parent->KeyAt(0)) < 0){
                new_node->SetParentPageId(parent->GetPageId());
            }else if(comparator_(key, new_parent->KeyAt(0)) == 0){
                new_node->SetParentPageId(new_parent->GetPageId());
            }else{
                new_node->SetParentPageId(new_parent->GetPageId());
                old_node->SetParentPageId(new_parent->GetPageId());
            }
            buffer_pool_manager_->UnpinPage(new_node->GetPageId(), true);
            // delete tmp page
            buffer_pool_manager_->UnpinPage(tmp_node->GetPageId(), false);
            buffer_pool_manager_->DeletePage(tmp_node->GetPageId());
            // insert parent's parent
            InsertIntoParent(parent, new_parent->KeyAt(0), new_parent);
        }
        buffer_pool_manager_->UnpinPage(parent->GetPageId(), true);
    }
}

/*!
 *
 * @tparam KeyType
 * @tparam ValueType
 * @tparam KeyComparator
 * @param key
 * @param result
 * @param transaction
 * @return
 */
template <typename KeyType, typename ValueType, typename KeyComparator>
bool BPlusTree<KeyType, ValueType, KeyComparator>::
GetValue(const KeyType &key, std::vector<ValueType> &result,
         Transaction *transaction){
    // step 1: find leaf node
    auto leaf = FindLeafPage(key, false, Operation::READONLY, transaction);
    if(leaf == nullptr){
        return false;
    }else{
        bool ret = false;
        ValueType value;
        if(leaf->Lookup(key, value, comparator_)){
            result.push_back(value);
            ret = true;
        }
        if(transaction == nullptr){
            auto page_id = leaf->GetPageId();
            buffer_pool_manager_->FetchPage(page_id)->RUnlatch();
            buffer_pool_manager_->UnpinPage(page_id, false);
            buffer_pool_manager_->UnpinPage(page_id, false);
        }else{
            UnlockUnpinPages(Operation::READONLY, transaction);
        }
        return ret;
    }
}

